import cv2
import numpy as np
from numpy import ndarray 
from enum import Enum, auto

# 定义攻击类型枚举
class AttackType(Enum):
    NONE = auto()  # 无攻击
    JPEG_COMPRESSION = auto()  # JPEG压缩
    GAUSSIAN_FILTER = auto()  # 高斯低通滤波
    HISTOGRAM_EQUALIZATION = auto()  # 直方图均衡化
    BRIGHTEN = auto()  # 图像增亮
    DARKEN = auto()  # 图像变暗
    INCREASE_CONTRAST = auto()  # 增加对比度
    DECREASE_CONTRAST = auto()  # 降低对比度
    ADD_GAUSSIAN_NOISE = auto()  # 添加高斯噪声
    ADD_SALT_PEPPER_NOISE = auto()  # 添加椒盐噪声
    ADD_MULTIPLICATIVE_NOISE = auto()  # 添加乘积性噪声
    MEDIAN_FILTER = auto()  # 中值滤波
    CROP = auto()  # 剪切
    ROTATE = auto()  # 旋转
    SCALE = auto()  # 缩放


class Attacker:

    AttackType = AttackType

# 攻击函数
    @staticmethod
    def attack(original_image:ndarray, attack_type: AttackType):
        image = original_image.copy()
        if attack_type == AttackType.NONE:
            return image
        elif attack_type == AttackType.JPEG_COMPRESSION:
            _, img_encoded = cv2.imencode('.jpg', image, [int(cv2.IMWRITE_JPEG_QUALITY), 50])
            return cv2.imdecode(img_encoded, cv2.IMREAD_GRAYSCALE)
        elif attack_type == AttackType.GAUSSIAN_FILTER:
            return cv2.GaussianBlur(image, (3, 3), 1)
        elif attack_type == AttackType.HISTOGRAM_EQUALIZATION:
            return cv2.equalizeHist(image)
        elif attack_type == AttackType.BRIGHTEN:
            return cv2.convertScaleAbs(image, alpha=1.2, beta=50)
        elif attack_type == AttackType.DARKEN:
            return cv2.convertScaleAbs(image, alpha=0.5, beta=-50)
        elif attack_type == AttackType.INCREASE_CONTRAST:
            return cv2.convertScaleAbs(image, alpha=1.5, beta=0)
        elif attack_type == AttackType.DECREASE_CONTRAST:
            return cv2.convertScaleAbs(image, alpha=0.5, beta=0)
        elif attack_type == AttackType.ADD_GAUSSIAN_NOISE:
            noise = np.random.normal(0, 0.05, image.shape)
            return cv2.add(image, noise.astype(np.float64))
        elif attack_type == AttackType.ADD_SALT_PEPPER_NOISE:
            salt_prob = 0.01
            pepper_prob = 0.01
            noisy_image = np.copy(image)
            total_pixels = image.size
            num_salt = int(salt_prob * total_pixels)
            num_pepper = int(pepper_prob * total_pixels)

            # 添加椒噪声
            coords = [np.random.randint(0, i - 1, num_salt) for i in image.shape]
            noisy_image[coords] = 255

            # 添加盐噪声
            coords = [np.random.randint(0, i - 1, num_pepper) for i in image.shape]
            noisy_image[coords] = 0

            return noisy_image
        elif attack_type == AttackType.ADD_MULTIPLICATIVE_NOISE:
            noise = np.random.normal(0, 0.01, image.shape)
            return cv2.add(image, image * noise.astype(np.float32))
        elif attack_type == AttackType.MEDIAN_FILTER:
            return cv2.medianBlur(image, 5)
        elif attack_type == AttackType.CROP:
            image[200:256, 200:256] = 244
            return image
        elif attack_type == AttackType.ROTATE:
            rows, cols = image.shape
            M = cv2.getRotationMatrix2D((cols / 2, rows / 2), 10, 1)
            return cv2.warpAffine(image, M, (cols, rows))
        elif attack_type == AttackType.SCALE:
            resized_image = cv2.resize(image, (256, 256))
            return cv2.resize(resized_image, (512, 512))
        else:
            raise ValueError("未知的攻击类型")